package vn.edu.hcmut.cse.ai.search.local;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.List;
import vn.edu.hcmut.cse.ai.infrastructure.Action;
import vn.edu.hcmut.cse.ai.infrastructure.Node;
import vn.edu.hcmut.cse.ai.infrastructure.Problem;
import vn.edu.hcmut.cse.ai.infrastructure.Solution;
import vn.edu.hcmut.cse.ai.search.SearchAlgorithm;

public class SimpleHillClimbingSearch implements SearchAlgorithm {

    private final HeuristicFunction heuristicFunction;

    public SimpleHillClimbingSearch(HeuristicFunction function) {
        this.heuristicFunction = function;
    }

    @Override
    public List<Action> doSearching(Problem problem) {
        Node rootNode = new Node(problem.getInitialState(), null, null);
        double initialDistance = this.heuristicFunction.heuristicEvaluation(problem.getInitialState(), problem.getGoalTest().getGoalState());
        rootNode.setHeuristicPoint(initialDistance);

        // Add a hash set to reduce the expanded node list
        HashSet<Object> exploredSet = new HashSet<>();
        exploredSet.add(rootNode.getState());

        Node currentNode = rootNode;
        boolean canFindBetterNode;
        while (true) {
            canFindBetterNode = false;
            if (problem.getGoalTest().isGoalState(currentNode.getState())) {
                return new Solution(currentNode).getActionList();
            }

            // Find the expanded nodes of the current node
            List<Node> expandedNodes = expandNode(currentNode, problem, exploredSet);

            // Loop for all expanded nodes to find a better node
            for (int i = 0; i < expandedNodes.size(); i++) {
                Node childNode = expandedNodes.get(i);
                childNode.setHeuristicPoint(this.heuristicFunction.heuristicEvaluation(childNode.getState(), problem.getGoalTest().getGoalState()));

                if (currentNode.getHeuristicPoint() < childNode.getHeuristicPoint()) {
                    canFindBetterNode = true;
                    currentNode = childNode;
                    exploredSet.add(currentNode.getState());
                    break;
                }
            }

            if (canFindBetterNode == false) {
                break;
            }
        }
        return new ArrayList<>();
    }

    private List<Node> expandNode(Node node, Problem problem, HashSet<Object> exploredSet) {

        List<Action> actionList = problem.getActions().getActionList(node.getState());
        List<Node> expandedNodeList = new ArrayList<>();

        for (Action anAction : actionList) {
            Node childNode = problem.getChildNode(node, anAction);
            if (!exploredSet.contains(childNode.getState())) {
                expandedNodeList.add(childNode);
            }
        }

        // Shuffle the list to make it random
        Collections.shuffle(expandedNodeList);

        return expandedNodeList;
    }

}
